"""
ChromaDB向量存储实现
支持高效的语义搜索和文档检索
"""

import asyncio
from typing import List, Dict, Any, Optional
import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions
import numpy as np
from pathlib import Path

from ...core.interfaces import VectorStoreInterface, Document, SearchResult
from ...core.events import EventEmitter, SystemEvents
from ...core.config import config


class ChromaVectorStore(VectorStoreInterface, EventEmitter):
    """ChromaDB向量存储实现"""
    
    def __init__(self, persist_directory: str = "./data/vector_db", collection_name: str = "hr_documents"):
        super().__init__()
        self.persist_directory = Path(persist_directory)
        self.collection_name = collection_name
        self.client = None
        self.collection = None
        self._embedding_function = None
        
    async def initialize(self):
        """初始化ChromaDB客户端"""
        try:
            # 确保目录存在
            self.persist_directory.mkdir(parents=True, exist_ok=True)
            
            # 创建ChromaDB客户端
            self.client = chromadb.PersistentClient(
                path=str(self.persist_directory),
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            
            # 设置嵌入函数
            embedding_model = config.get("lightrag_engine.embedding.model", "BAAI/bge-base-zh-v1.5")
            self._embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
                model_name=embedding_model
            )
            
            # 获取或创建集合
            try:
                self.collection = self.client.get_collection(
                    name=self.collection_name,
                    embedding_function=self._embedding_function
                )
            except Exception:
                self.collection = self.client.create_collection(
                    name=self.collection_name,
                    embedding_function=self._embedding_function,
                    metadata={"hnsw:space": "cosine"}
                )
            
            await self.emit_success("vector_store_initialized", {
                "collection_name": self.collection_name,
                "persist_directory": str(self.persist_directory)
            })
            
        except Exception as e:
            await self.emit_error("vector_store_initialization", e)
            raise
    
    async def add_documents(self, documents: List[Document]) -> bool:
        """添加文档到向量存储"""
        if not self.collection:
            await self.initialize()
        
        try:
            # 准备数据
            ids = [doc.id for doc in documents]
            documents_text = [doc.content for doc in documents]
            metadatas = [doc.metadata for doc in documents]
            
            # 如果文档已有嵌入向量，使用它们；否则让ChromaDB自动生成
            embeddings = None
            if all(doc.embedding for doc in documents):
                embeddings = [doc.embedding for doc in documents]
            
            # 批量添加文档
            if embeddings:
                self.collection.add(
                    ids=ids,
                    documents=documents_text,
                    metadatas=metadatas,
                    embeddings=embeddings
                )
            else:
                self.collection.add(
                    ids=ids,
                    documents=documents_text,
                    metadatas=metadatas
                )
            
            await self.emit(SystemEvents.DOCUMENT_INDEXED, {
                "count": len(documents),
                "collection": self.collection_name
            })
            
            return True
            
        except Exception as e:
            await self.emit_error("add_documents", e)
            return False
    
    async def search_similar(
        self, 
        query_embedding: List[float], 
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """向量相似度搜索"""
        if not self.collection:
            await self.initialize()
        
        try:
            # 构建查询参数
            query_params = {
                "query_embeddings": [query_embedding],
                "n_results": top_k,
                "include": ["documents", "metadatas", "distances"]
            }
            
            # 添加元数据过滤
            if filter_metadata:
                query_params["where"] = filter_metadata
            
            # 执行查询
            results = self.collection.query(**query_params)
            
            # 转换结果
            search_results = []
            for i in range(len(results["ids"][0])):
                doc = Document(
                    id=results["ids"][0][i],
                    content=results["documents"][0][i],
                    metadata=results["metadatas"][0][i] or {}
                )
                
                # ChromaDB返回的是距离，需要转换为相似度分数
                distance = results["distances"][0][i]
                similarity_score = 1.0 / (1.0 + distance)  # 转换为0-1之间的相似度
                
                search_results.append(SearchResult(
                    document=doc,
                    score=similarity_score,
                    relevance_type="vector"
                ))
            
            await self.emit(SystemEvents.SEARCH_PERFORMED, {
                "query_type": "vector_similarity",
                "results_count": len(search_results),
                "top_k": top_k
            })
            
            return search_results
            
        except Exception as e:
            await self.emit_error("search_similar", e)
            return []
    
    async def search_by_text(
        self, 
        query_text: str, 
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """文本查询搜索"""
        if not self.collection:
            await self.initialize()
        
        try:
            # 构建查询参数
            query_params = {
                "query_texts": [query_text],
                "n_results": top_k,
                "include": ["documents", "metadatas", "distances"]
            }
            
            # 添加元数据过滤
            if filter_metadata:
                query_params["where"] = filter_metadata
            
            # 执行查询
            results = self.collection.query(**query_params)
            
            # 转换结果（与search_similar相同的逻辑）
            search_results = []
            for i in range(len(results["ids"][0])):
                doc = Document(
                    id=results["ids"][0][i],
                    content=results["documents"][0][i],
                    metadata=results["metadatas"][0][i] or {}
                )
                
                distance = results["distances"][0][i]
                similarity_score = 1.0 / (1.0 + distance)
                
                search_results.append(SearchResult(
                    document=doc,
                    score=similarity_score,
                    relevance_type="vector"
                ))
            
            return search_results
            
        except Exception as e:
            await self.emit_error("search_by_text", e)
            return []
    
    async def update_document(self, doc_id: str, document: Document) -> bool:
        """更新文档"""
        if not self.collection:
            await self.initialize()
        
        try:
            # ChromaDB的更新方式是先删除再添加
            await self.delete_document(doc_id)
            await self.add_documents([document])
            
            await self.emit(SystemEvents.DOCUMENT_UPDATED, {
                "document_id": doc_id,
                "collection": self.collection_name
            })
            
            return True
            
        except Exception as e:
            await self.emit_error("update_document", e)
            return False
    
    async def delete_document(self, doc_id: str) -> bool:
        """删除文档"""
        if not self.collection:
            await self.initialize()
        
        try:
            self.collection.delete(ids=[doc_id])
            
            await self.emit(SystemEvents.DOCUMENT_DELETED, {
                "document_id": doc_id,
                "collection": self.collection_name
            })
            
            return True
            
        except Exception as e:
            await self.emit_error("delete_document", e)
            return False
    
    async def get_collection_stats(self) -> Dict[str, Any]:
        """获取集合统计信息"""
        if not self.collection:
            await self.initialize()
        
        try:
            count = self.collection.count()
            return {
                "collection_name": self.collection_name,
                "document_count": count,
                "persist_directory": str(self.persist_directory)
            }
        except Exception as e:
            await self.emit_error("get_collection_stats", e)
            return {}
    
    async def clear_collection(self) -> bool:
        """清空集合"""
        if not self.collection:
            await self.initialize()
        
        try:
            # 获取所有文档ID
            results = self.collection.get()
            if results["ids"]:
                self.collection.delete(ids=results["ids"])
            
            await self.emit("vector_collection_cleared", {
                "collection_name": self.collection_name
            })
            
            return True
            
        except Exception as e:
            await self.emit_error("clear_collection", e)
            return False